# 调试/实验脚本
"""
代码尝试专用脚本（空壳模板）
- 用于调试、测试新功能或片段
使用 rich 替换终端输出与进度条。
"""
import time
from rich.progress import Progress, SpinnerColumn, BarColumn, TextColumn

from utils.plot_style import print_start, print_end, set_chinese_font, set_plot_style, force_chinese_font

from src.plot import prepare_border_counties_population_data
import matplotlib.pyplot as plt
import matplotlib as mpl
from matplotlib.ticker import FuncFormatter
from pathlib import Path
import io
import imageio
import numpy as np


def run_tests():
    """在这里放置你的快速试验代码或单元级小片段。"""
    from rich.console import Console
    console = Console()
    console.rule("try.py: 运行示例测试")
    # 非破坏性测试：尝试读取 config 指定的数据文件（若不存在则打印提示）
    try:
        from src.core_utils import load_data
        try:
            df = load_data()
            console.log(f"load_data 成功，rows={len(df)} cols={len(df.columns)}")
            # 仅示例性加载数据并打印行列信息（非破坏性）
            console.log("注意：回归与示例已被移除，请参见 examples/ 目录中的示例脚本。")
        except FileNotFoundError as e:
            console.log(f"load_data 未找到文件：{e}")
        except Exception as e:
            console.log(f"load_data 出错：{e}")
    except Exception:
        console.log("src.core_utils.load_data 不可用：跳过数据加载测试")
    # 演示 rich 进度条
    with Progress(
        SpinnerColumn(style="red"),
        TextColumn("[bold]测试进度"),
        BarColumn(bar_width=None, complete_style="red"),
        TextColumn("{task.completed}/{task.total}"),
    ) as progress:
        task = progress.add_task("test", total=5)
        for _ in range(5):
            time.sleep(0.2)
            progress.advance(task)


def generate_heatmap_from_prepared_data(year: int = 2020, out_dir: str = 'figures', scope: str = 'county') -> str:
    """使用 src.plot.prepare_border_counties_population_data 准备数据并在此处完成绘图与保存。

    返回保存的相对路径字符串。
    """
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    # 准备数据（不在 src.plot 中绘图/保存）
    prepared = prepare_border_counties_population_data(year=year, scope=scope)
    geo = prepared['geo']
    cmap = prepared['cmap']
    vmin = prepared['vmin']
    vmax = prepared['vmax']
    plot_figsize = prepared['figsize']
    plot_dpi = prepared['dpi']
    filename = prepared['filename']
    # filename 由 prepare 返回（已包含中文命名），无需在此追加后缀
    cb_bbox = prepared['colorbar_bbox']
    font_family = prepared['font_family']

    # 应用样式配置（如果 prepare 返回了 cfg_plot，则使用它以保持一致性）
    plot_cfg = prepared.get('cfg_plot') or {}
    set_plot_style(plot_cfg or {})

    # 仅在字体可用时加载中文字体，避免重复或不可用字体导致的警告
    if prepared.get('font_available') and prepared.get('font_family'):
        set_chinese_font(prepared.get('font_family'))
    else:
        # 建议字体不可用或未指定：静默回退到系统默认字体（避免在终端生成多余提示）
        pass

    fig, ax = plt.subplots(figsize=plot_figsize, dpi=plot_dpi)
    geo.plot(column='总人口', cmap=cmap, linewidth=0.4, edgecolor='#444444', ax=ax)
    ax.set_axis_off()

    mappable = mpl.cm.ScalarMappable(norm=mpl.colors.Normalize(vmin=vmin, vmax=vmax), cmap=cmap)
    mappable._A = []
    cax = fig.add_axes(cb_bbox)
    cb = fig.colorbar(mappable, cax=cax, orientation='horizontal')

    def per_ten_thousand(x, pos):
        val = x / 10000.0
        if val >= 10:
            return f"{val:.0f}"
        return f"{val:.1f}"

    cb.ax.xaxis.set_major_formatter(FuncFormatter(per_ten_thousand))
    cb.ax.tick_params(labelsize=9)
    cb.set_label('七普边境县乡镇总人口数（万人）', fontsize=11)

    force_chinese_font(ax, font_family)
    for label in cb.ax.get_xticklabels():
        label.set_fontfamily(font_family)
    cb.ax.xaxis.label.set_fontfamily(font_family)

    fig.subplots_adjust(top=0.92, left=0.06, right=0.98, bottom=0.06)
    out_path = out_dir / filename
    bbox_opt = (prepared['cfg_plot'].get('bbox_inches') if prepared.get('cfg_plot') and prepared['cfg_plot'].get('bbox_inches') else 'tight')
    fig.savefig(out_path, dpi=plot_dpi, bbox_inches=bbox_opt)
    repo_root = Path(__file__).resolve().parents[1]
    try:
        rel_path = out_path.relative_to(repo_root)
    except Exception:
        rel_path = out_path
    from rich.console import Console
    Console().log(f'已保存热力图到 {rel_path}')
    return str(rel_path)


def generate_density_gif(years, out_path: str = 'figures/边境乡镇_mean_density_2000-2020.gif', column: str = 'mean', cmap: str = 'Reds', dpi: int = 150, figsize=(10, 12), scope: str = 'county', duration: float = 0.6, population_parquet: str = 'data/边境城市乡镇人口密度面板数据.parquet'):
    """为给定年份列表渲染每年人口密度字段（默认为 'mean'）并合成 GIF。

    说明：使用 prepare_border_counties_population_data 来准备每年的 GeoDataFrame，
    但不保存中间静态图像，直接把 matplotlib 图转为 RGBA 数组并传给 imageio。
    """
    out_dir = Path('figures')
    out_dir.mkdir(parents=True, exist_ok=True)
    frames = []

    # 为了保证色标一致性，先计算全时间段的 vmin/vmax
    vmin = None
    vmax = None
    geos = {}
    for y in years:
        prep = prepare_border_counties_population_data(year=y, scope=scope, population_parquet=population_parquet)
        if not prep:
            continue
        geos[y] = prep
        if vmin is None or prep['vmin'] < vmin:
            vmin = prep['vmin']
        if vmax is None or prep['vmax'] > vmax:
            vmax = prep['vmax']

    if not geos:
        from rich.console import Console
        Console().log('未找到任何年份的数据，GIF 生成终止。')
        return None

    # 渲染每一年为一帧
    for y in sorted(geos.keys()):
        prep = geos[y]
        geo = prep['geo']
        plot_figsize = prep['figsize']
        plot_dpi = prep['dpi']
        font_family = prep.get('font_family')

        # 样式
        set_plot_style(prep.get('cfg_plot') or {})
        if prep.get('font_available') and font_family:
            set_chinese_font(font_family)

        fig, ax = plt.subplots(figsize=figsize or plot_figsize, dpi=plot_dpi)
        # 绘制指定列作为热力图
        if column not in geo.columns:
            # 如果没有该列，尝试使用 'mean' 或 '总人口' 作为回退
            col = 'mean' if 'mean' in geo.columns else '总人口'
        else:
            col = column

        geo.plot(column=col, cmap=cmap, linewidth=0.2, edgecolor='#444444', ax=ax, vmin=vmin, vmax=vmax)
        ax.set_axis_off()

        mappable = mpl.cm.ScalarMappable(norm=mpl.colors.Normalize(vmin=vmin, vmax=vmax), cmap=cmap)
        mappable._A = []
        cax = fig.add_axes(prep['colorbar_bbox'])
        cb = fig.colorbar(mappable, cax=cax, orientation='horizontal')

        def per_ten_thousand(x, pos):
            val = x / 10000.0
            if val >= 10:
                return f"{val:.0f}"
            return f"{val:.1f}"

        try:
            cb.ax.xaxis.set_major_formatter(FuncFormatter(per_ten_thousand))
            cb.ax.tick_params(labelsize=9)
            cb.set_label('人口密度（单位根据数据）', fontsize=11)
        except Exception:
            pass

        force_chinese_font(ax, font_family)

        # 在内存中保存为 RGBA 图像并加入 frames
        buf = io.BytesIO()
        fig.savefig(buf, format='png', dpi=plot_dpi, bbox_inches='tight')
        plt.close(fig)
        buf.seek(0)
        img = imageio.v2.imread(buf)
        frames.append(img)

    # 写入 GIF
    gif_path = Path(out_path)
    imageio.mimsave(gif_path, frames, format='GIF', duration=duration)
    from rich.console import Console
    try:
        repo_root = Path(__file__).resolve().parents[1]
        rel = gif_path.relative_to(repo_root)
    except Exception:
        rel = gif_path
    Console().log(f'已保存 GIF 到 {rel} (frames={len(frames)})')
    return str(rel)


if __name__ == "__main__":
    start_time = print_start('生成热力图（try.py）')
    # 直接生成热力图并保存
    try:
        # 生成边境县范围热力图（原先行为）
        generate_heatmap_from_prepared_data(year=2020, out_dir='figures', scope='county')
        # 生成边境城市范围热力图（新的需求）
        generate_heatmap_from_prepared_data(year=2020, out_dir='figures', scope='city')
    except Exception as e:
        from rich.console import Console
        Console().log(f'生成热力图出错：{e}')
    print_end(start_time, '生成热力图（try.py）')

